import copy
from path import PathMgr
from res import ResMgr
from reader import IntReader
from config import Config

# 初赛：输入一个统一的接受站点功耗                          [OK]
# 复赛：不输入统一的功耗，增加一行输入每个的卫星功耗        [OK]
# 决赛：增加一行输入卫星容量和地面站点用户数，增加容量限制  [OK]

class Case:
    def GetNodePair(send, recv):
        if send < recv:
            return send * Config.NODE_MAX_ID + recv
        else:
            return recv * Config.NODE_MAX_ID + send

    def __init__(self, stageId, caseId):
        self.caseId = caseId
        self.ReadConfig(stageId, caseId)
        return

    def IsPreStage(self):
        return self.__stageId == Config.STAGE_PRE_ID
    def IsMidStage(self):
        return self.__stageId == Config.STAGE_MID_ID
    def IsFinalStage(self):
        return (self.__stageId == Config.STAGE_FINAL_ID) or (self.__stageId == Config.STAGE_TEST_ID)

    def AddEdge(self, send, recv, dist):
        self.__edgeMgr[Case.GetNodePair(send, recv)] = dist
        return
    def DelEdge(self, send, recv):
        self.__edgeMgr.pop(Case.GetNodePair(send, recv), -1)
        return
    def GetDist(self, send, recv):
        return self.__edgeMgr.get(Case.GetNodePair(send, recv), -1)

    def IsSiteIdValid(self, siteId):
        return (siteId >= 0) and (siteId < self.__siteNum)
    
    def IsRouteDistValid(self, routeDist):
        return (routeDist >=0) and (routeDist <= self.__maxDist)

    def GetPwr(self, siteId):
        if self.IsPreStage():
            return self.__recvPwr
        else:
            return self.__sitePwrMgr[siteId]

    def DecUserCnt(self, recvSiteId, groundSiteId):
        if not self.IsFinalStage():
            return 0
        curCnt = self.__siteUserCntMgr[recvSiteId]
        decCnt = self.__siteUserCntMgr[groundSiteId]
        if curCnt < decCnt:
            return (decCnt - curCnt)
        self.__siteUserCntMgr[recvSiteId] = curCnt - decCnt
        return 0

    def ReadConfig(self, stageId, caseId):
        intReader = IntReader(PathMgr.GetCasePath(caseId))
        # 固定数据
        self.__stageId = stageId
        self.__siteNum = intReader.Read()
        self.__edgeNum = intReader.Read()
        self.castPwr = intReader.Read()
        self.__maxDist = intReader.Read()

        if self.IsPreStage():
            self.__recvPwr = intReader.Read()

        self.__sitePwrMgr = []

        # 动态数据
        self.__siteUserCntMgr = []
        self.__edgeMgr = {}
        self.groundSiteSet = set()
        self.spaceSiteSet = set()

        for i in range(self.__siteNum):
            siteType = intReader.Read()
            if siteType == 0:
                self.groundSiteSet.add(i)
            elif siteType == 1:
                self.spaceSiteSet.add(i)

        if not self.IsPreStage():
            for i in range(self.__siteNum):
                self.__sitePwrMgr.append(intReader.Read())
        if self.IsFinalStage():
            for i in range(self.__siteNum):
                self.__siteUserCntMgr.append(intReader.Read())
        for i in range(self.__edgeNum):
            send = intReader.Read()
            recv = intReader.Read()
            dist = intReader.Read()
            self.AddEdge(send, recv, dist)
        return

class CaseMgr:
    def Init():
        CaseMgr.__caseVec = []
        CaseMgr.LoadCases()
        return

    def LoadCases():
        for caseId in range(Config.CASE_NUM):
            CaseMgr.__caseVec.append(Case(Config.STAGE_ID, caseId + 1))
        return

    def GetCommonCase(caseId):
        return CaseMgr.__caseVec[caseId - 1]
    def GetCloneCase(caseId):
        case = copy.deepcopy(CaseMgr.__caseVec[caseId - 1])
        return case